import math
import threading
import zlib

from Crypto.Cipher import AES
import ipaddr
import struct
import time
import os
import zstandard as zstd

decompress_lock = threading.Lock()


class Pcap:
    internals = {
        'pr2name': {
            1: 'icmp',
            2: 'igmp',
            6: 'tcp',
            17: 'udp',
            47: 'gre',
            50: 'esp',
            58: 'icmpv6',
            89: 'ospf',
            103: 'pim',
            132: 'sctp'
        },
        'pcaps': {}
    }
    shortHeader = None

    def __init__(self, key):
        self.iv = None
        self._count = 0
        self._closing = False
        self.key = key
        self.fd = None
        self.corrupt = False
        self.headBuffer = None

    def decompressStreaming(self, compressed_data, form):
        try:
            with decompress_lock:
                with zstd.ZstdDecompressor().stream_reader(compressed_data) as reader:
                    decompressed_data = reader.read()
                    return decompressed_data
        except Exception as e:
            print(e)
            return bytearray()

    @classmethod
    def get(cls, key):
        if key in cls.internals['pcaps']:
            return cls.internals['pcaps'][key]
        pcap = cls(key)
        cls.internals['pcaps'][key] = pcap
        return pcap

    def is_open(self):
        return self.fd is not None

    def is_corrupt(self):
        return self.corrupt

    def open(self, info):
        if self.fd:
            return
        self.filename = info['name']
        self.encoding = info.get('encoding', 'normal')

        if 'dek' in info:
            decipher = AES.new(info['kek'].encode(), AES.MODE_CBC)
            dek = bytes.fromhex(info['dek'])
            self.encKey = decipher.decrypt(dek)

        if 'uncompressedBits' in info:
            self.uncompressedBits = info['uncompressedBits']
            self.uncompressedBitsSize = 2 ** self.uncompressedBits
            self.compression = 'gzip'

        if 'compression' in info:
            self.compression = info['compression']

        if 'iv' in info:
            iv = bytes.fromhex(info['iv'])
            self.iv = bytearray(16)
            self.iv[:len(iv)] = iv

        try:
            self.fd = os.open(info['name'], os.O_RDONLY)
            # self.read_header()
        except Exception as err:
            self.fd = None
            raise err

    def ref(self):
        self._count += 1

    def unref(self):
        self._count -= 1
        if self._count > 0:
            return

        if self._closing:
            return
        self._closing = True

        def close_fd():
            if self._closing and self._count == 0:
                del Pcap.internals['pcaps'][self.key]
                if self.fd:
                    os.close(self.fd)
                self.fd = None
            else:
                self._closing = False

        time.sleep(0.5)
        close_fd()

    def create_decipher(self, pos):
        self.iv[12:16] = struct.pack('>I', pos)
        return AES.new(self.encKey, AES.MODE_CTR, nonce=self.iv)

    def read_header(self):
        if self.headBuffer:
            return self.headBuffer
        # pcap header is 24 but reading extra because of gzip/encryption
        self.headBuffer = bytearray(64)
        # os.lseek(self.fd, 0, os.SEEK_SET)
        self.headBuffer = os.read(self.fd, len(self.headBuffer))
        if self.encoding == 'aes-256-ctr':
            decipher = self.create_decipher(0)
            self.headBuffer = bytearray(decipher.decrypt(bytes(self.headBuffer)))
        elif self.encoding == 'xor-2048':
            for i in range(len(self.headBuffer)):
                self.headBuffer[i] ^= self.encKey[i % 256]

        if self.uncompressedBits:
            if self.compression == 'gzip':
                self.headBuffer = zlib.decompress(bytes(self.headBuffer), zlib.MAX_WBITS | 16)
            elif self.compression == 'zstd':
                self.headBuffer = self.decompressStreaming(bytes(self.headBuffer), 2)

        self.headBuffer = self.headBuffer[:24]

        magic = struct.unpack('<I', self.headBuffer[:4])[0]
        self.bigEndian = (magic == 0xd4c3b2a1 or magic == 0x4d3cb2a1)
        self.nanosecond = (magic == 0xa1b23c4d or magic == 0x4d3cb2a1)

        if not self.bigEndian and magic not in {0xa1b2c3d4, 0xa1b23c4d, 0xa1b2c3d5}:
            self.corrupt = True
            os.close(self.fd)
            raise ValueError("Corrupt PCAP header")

        if magic == 0xa1b2c3d5:
            self.shortHeader = struct.unpack('<I', self.headBuffer[8:12])[0]
            self.headBuffer[0] = 0xd4  # Reset header to normal

        self.linkType = struct.unpack('>I' if self.bigEndian else '<I', self.headBuffer[20:24])[0]
        return self.headBuffer

    def read_packet(self, pos):
        if not self.fd:
            time.sleep(0.01)
            return self.read_packet(pos)
        return self.read_packet_internal(pos, -1)

    def read_packet_internal(self, pos_arg, hp_len_arg):
        pos = pos_arg
        hp_len = hp_len_arg

        if hp_len == -1:
            if self.compression == "zstd":
                hp_len = self.uncompressedBitsSize
            else:
                hp_len = 2048

        inside_offset = 0
        if self.uncompressedBits:
            inside_offset = pos & self.uncompressedBitsSize - 1
            pos = math.floor(pos / self.uncompressedBitsSize)

        pos_offset = 0

        if self.encoding == 'aes-256-ctr':
            pos_offset = pos % 16
            pos = pos - pos_offset
        elif self.encoding == 'xor-2048':
            pos_offset = pos % 256
            pos = pos - pos_offset

        hp_len = 256 * math.ceil((hp_len + inside_offset + pos_offset) / 256)
        buffer = bytearray(hp_len)
        os.lseek(self.fd, pos, os.SEEK_SET)
        read_buffer = os.read(self.fd, len(buffer))
        # read_buffer = buffer[:bytes_read]
        if len(read_buffer) - pos_offset < 16:
            return None
        if self.encoding == 'aes-256-ctr':
            decipher = self.create_decipher(pos // 16)
            read_buffer = bytearray(decipher.decrypt(read_buffer))[pos_offset:]
        elif self.encoding == 'xor-2048':
            read_buffer = bytearray(b ^ self.encKey[i % 256] for i, b in enumerate(read_buffer))[pos_offset:]

        if self.uncompressedBits:
            try:
                if self.compression == 'gzip':
                    read_buffer = zlib.decompress(read_buffer, zlib.MAX_WBITS | 16)
                elif self.compression == 'zstd':
                    read_buffer = self.decompressStreaming(read_buffer, 1)
            except Exception as e:
                print(f"PCAP uncompress issue: {self.key} {pos} {len(buffer)} {read_buffer} {e}")
                return None

        if inside_offset:
            read_buffer = read_buffer[inside_offset:]

        header_len = 16 if self.shortHeader is None else 6

        if len(read_buffer) < header_len:
            if hp_len_arg == -1 and self.compression == 'zstd':
                return self.read_packet_internal(pos_arg, self.uncompressedBitsSize * 2)
            print(f"Not enough data {len(read_buffer)} for header {header_len}")
            return None

        packet_len = struct.unpack('>I' if self.bigEndian else '<I', read_buffer[8:12])[
            0] if self.shortHeader is None else struct.unpack('>H' if self.bigEndian else '<H', read_buffer[:2])[0]

        if packet_len < 0 or packet_len > 0xffff:
            return None

        if header_len + packet_len <= len(read_buffer):
            if self.shortHeader is not None:
                t = struct.unpack('<I', read_buffer[2:6])[0]
                sec = (t >> 20) + self.shortHeader
                usec = t & 0xfffff

                new_buffer = bytearray(16 + packet_len)
                struct.pack_into('<I', new_buffer, 0, sec)
                struct.pack_into('<I', new_buffer, 4, usec)
                struct.pack_into('<I', new_buffer, 8, packet_len)
                struct.pack_into('<I', new_buffer, 12, packet_len)
                new_buffer[16:] = read_buffer[6:packet_len + 6]
                return bytes(new_buffer)
            return read_buffer[:header_len + packet_len]

        if hp_len_arg != -1:
            return None

        return self.read_packet_internal(pos_arg, 16 + packet_len)

    @staticmethod
    def inet_ntoa(num):
        return f"{(num >> 24) & 0xff}.{(num >> 16) & 0xff}.{(num >> 8) & 0xff}.{num & 0xff}"

    def icmp(self, buffer, obj, pos):
        obj['icmp'] = {
            '_pos': pos,
            'length': len(buffer),
            'type': buffer[0],
            'code': buffer[1],
            'sum': int.from_bytes(buffer[2:4], byteorder='big')
        }
        obj['icmp']['data'] = buffer

    def tcp(self, buffer, obj, pos):
        try:
            obj['tcp'] = {
                '_pos': pos,
                'length': len(buffer),
                'sport': int.from_bytes(buffer[0:2], byteorder='big'),
                'dport': int.from_bytes(buffer[2:4], byteorder='big'),
                'seq': int.from_bytes(buffer[4:8], byteorder='big'),
                'ack': int.from_bytes(buffer[8:12], byteorder='big'),
                'off': (buffer[12] >> 4) & 0xf,
                'res1': buffer[12] & 0xf,
                'flags': buffer[13],
                'res2': (buffer[13] >> 6) & 0x3,
                'urgflag': (buffer[13] >> 5) & 0x1,
                'ackflag': (buffer[13] >> 4) & 0x1,
                'pshflag': (buffer[13] >> 3) & 0x1,
                'rstflag': (buffer[13] >> 2) & 0x1,
                'synflag': (buffer[13] >> 1) & 0x1,
                'finflag': buffer[13] & 0x1,
                'win': int.from_bytes(buffer[14:16], byteorder='big'),
                'sum': int.from_bytes(buffer[16:18], byteorder='big'),
                'urp': int.from_bytes(buffer[18:20], byteorder='big')
            }

            if 4 * obj['tcp']['off'] > len(buffer):
                obj['tcp']['data'] = bytes()
            else:
                obj['tcp']['data'] = buffer[4 * obj['tcp']['off']:]
        except Exception as e:
            print("Couldn't parse tcp", e)

    def udp(self, buffer, obj, pos):
        obj['udp'] = {
            '_pos': pos,
            'length': len(buffer),
            'sport': int.from_bytes(buffer[0:2], byteorder='big'),
            'dport': int.from_bytes(buffer[2:4], byteorder='big'),
            'ulen': int.from_bytes(buffer[4:6], byteorder='big'),
            'sum': int.from_bytes(buffer[6:8], byteorder='big')
        }

        obj['udp']['data'] = buffer[8:]
        data = obj['udp']['data']

        # vxlan
        if obj['udp']['dport'] == 4789 and len(data) > 8 and (data[0] & 0x77) == 0 and (data[1] & 0xb7) == 0:
            self.ether(buffer[16:], obj, pos + 16)

        # vxlan gpe
        if obj['udp']['dport'] == 4790 and len(data) > 8 and (data[0] & 0xf0) == 0 and (data[1] & 0xff) == 0:
            if data[3] == 1:
                self.ip4(buffer[16:], obj, pos + 16)
            elif data[3] == 2:
                self.ip6(buffer[16:], obj, pos + 16)
            elif data[3] == 3:
                self.ether(buffer[16:], obj, pos + 16)
            elif data[3] == 4:
                # TODO NSH
                pass

        # geneve
        if obj['udp']['dport'] == 6081 and len(data) > 8 and (data[0] & 0xc0) == 0 and (data[1] & 0x3f) == 0:
            optlen = data[0] & 0x3f
            protocol = (data[2] << 8) | data[3]
            offset = 8 + optlen * 4

            if 8 + offset < len(buffer):
                self.ethertyperun(protocol, buffer[8 + offset:], obj, pos + 8 + offset)

        # gtp
        if obj['udp']['dport'] == 2152 and len(data) > 8 and (data[0] & 0xf0) == 0x30 and data[1] == 0xff:
            offset = 8
            next_ext = 0
            if data[0] & 0x7:
                offset += 3
                next_ext = data[offset]
                offset += 1
            while next_ext != 0:
                extlen = data[offset]
                offset += 1
                offset += extlen * 4 - 2
                next_ext = data[offset]
                offset += 1

            if (data[offset] & 0xf0) == 0x60:
                self.ip6(data[offset:], obj, pos + offset)
            else:
                self.ip4(data[offset:], obj, pos + offset)

    def sctp(self, buffer, obj, pos):
        obj['sctp'] = {
            '_pos': pos,
            'length': len(buffer),
            'sport': int.from_bytes(buffer[0:2], byteorder='big'),
            'dport': int.from_bytes(buffer[2:4], byteorder='big')
        }
        obj['sctp']['data'] = buffer[12:]

    def esp(self, buffer, obj, pos):
        obj['esp'] = {
            '_pos': pos,
            'length': len(buffer)
        }
        obj['esp']['data'] = buffer

    def gre(self, buffer, obj, pos):
        obj['gre'] = {
            'flags_version': int.from_bytes(buffer[0:2], byteorder='big'),
            'type': int.from_bytes(buffer[2:4], byteorder='big')
        }

        bpos = 4
        if obj['gre']['flags_version'] & (0x8000 | 0x4000):
            bpos += 4

        # key
        if obj['gre']['flags_version'] & 0x2000:
            bpos += 4

        # sequence number
        if obj['gre']['flags_version'] & 0x1000:
            bpos += 4

        # routing
        if obj['gre']['flags_version'] & 0x4000:
            while True:
                bpos += 3
                length = int.from_bytes(buffer[bpos:bpos + 2], byteorder='big')
                bpos += 1
                if length == 0:
                    break
                bpos += length

        # ack number
        if obj['gre']['flags_version'] & 0x0080:
            bpos += 4

        if self.ethertyperun(obj['gre']['type'], buffer[bpos:], obj, pos + bpos):
            return

        if obj['gre']['type'] == 0x88be:
            self.ether(buffer[bpos + 8:], obj, pos + bpos + 8)
        else:
            print('gre Unknown type', obj['gre']['type'])

    def ip4(self, buffer, obj, pos):
        obj['ip'] = {
            'length': len(buffer),
            'hl': buffer[0] & 0xf,
            'v': (buffer[0] >> 4) & 0xf,
            'tos': buffer[1],
            'len': int.from_bytes(buffer[2:4], byteorder='big'),
            'id': int.from_bytes(buffer[4:6], byteorder='big'),
            'off': int.from_bytes(buffer[6:8], byteorder='big'),
            'ttl': buffer[8],
            'p': buffer[9],
            'sum': int.from_bytes(buffer[10:12], byteorder='big'),
            'addr1': self.inet_ntoa(int.from_bytes(buffer[12:16], byteorder='big')),
            'addr2': self.inet_ntoa(int.from_bytes(buffer[16:20], byteorder='big'))
        }

        payload = buffer[obj['ip']['hl'] * 4: obj['ip']['len']]
        pos_payload = pos + obj['ip']['hl'] * 4

        if obj['ip']['p'] == 1:
            self.icmp(payload, obj, pos_payload)
        elif obj['ip']['p'] == 6:
            self.tcp(payload, obj, pos_payload)
        elif obj['ip']['p'] == 17:
            self.udp(payload, obj, pos_payload)
        elif obj['ip']['p'] == 41:  # IPPROTO_IPV6
            self.ip6(payload, obj, pos_payload)
        elif obj['ip']['p'] == 50:  # IPPROTO_ESP
            self.esp(payload, obj, pos_payload)
        elif obj['ip']['p'] == 47:
            self.gre(payload, obj, pos_payload)
        elif obj['ip']['p'] == 132:
            self.sctp(payload, obj, pos_payload)
        else:
            obj['ip']['data'] = payload
            # print("v4 Unknown ip.p", obj)

    def ip6(self, buffer, obj, pos):
        obj['ip'] = {
            'length': len(buffer),
            'v': (buffer[0] >> 4) & 0xf,
            'tc': ((buffer[0] & 0xf) << 4) | ((buffer[1] >> 4) & 0xf),
            'flow': ((buffer[1] & 0xf) << 16) | (buffer[2] << 8) | buffer[3],
            'len': int.from_bytes(buffer[4:6], byteorder='big'),
            'p': buffer[6],
            'hopLimt': buffer[7],
            'addr1': ipaddr.IPv6Address(buffer[8:24]).exploded,
            'addr2': ipaddr.IPv6Address(buffer[24:40]).exploded
        }

        offset = 40
        while offset < len(buffer):
            if obj['ip']['p'] in {0, 60, 43}:  # IPPROTO_HOPOPTS, IPPROTO_DSTOPTS, IPPROTO_ROUTING
                obj['ip']['p'] = buffer[offset]
                offset += ((buffer[offset + 1] + 1) << 3)
            elif obj['ip']['p'] in {1, 58}:  # ICMPv6, ICMPv6 Router Advertisement
                self.icmp(buffer[offset:offset + obj['ip']['len']], obj, pos + offset)
                return
            elif obj['ip']['p'] == 4:  # IPPROTO_IPV4
                self.ip4(buffer[offset:offset + obj['ip']['len']], obj, pos + offset)
                return
            elif obj['ip']['p'] == 6:
                self.tcp(buffer[offset:offset + obj['ip']['len']], obj, pos + offset)
                return
            elif obj['ip']['p'] == 17:
                self.udp(buffer[offset:offset + obj['ip']['len']], obj, pos + offset)
                return
            elif obj['ip']['p'] == 47:
                self.gre(buffer[offset:offset + obj['ip']['len']], obj, pos + offset)
                return
            elif obj['ip']['p'] == 50:  # IPPROTO_ESP
                self.esp(buffer[offset:offset + obj['ip']['len']], obj, pos + offset)
                return
            elif obj['ip']['p'] == 132:
                self.sctp(buffer[offset:offset + obj['ip']['len']], obj, pos + offset)
                return
            else:
                obj['ip']['data'] = buffer[offset:offset + obj['ip']['len']]
                # print("v6 Unknown ip.p", obj)
                return

    def pppoe(self, buffer, obj, pos):
        obj['pppoe'] = {
            'len': int.from_bytes(buffer[4:6], byteorder='big') - 2,
            'type': int.from_bytes(buffer[6:8], byteorder='big')
        }

        if obj['pppoe']['type'] == 0x21:
            self.ip4(buffer[8:8 + obj['pppoe']['len']], obj, pos + 8)
        elif obj['pppoe']['type'] == 0x57:
            self.ip6(buffer[8:8 + obj['pppoe']['len']], obj, pos + 8)
        else:
            print('Unknown pppoe.type', obj)

    def ppp(self, buffer, obj, pos):
        obj['pppoe'] = {
            'type': int.from_bytes(buffer[2:4], byteorder='big')
        }

        if obj['pppoe']['type'] == 0x21:
            self.ip4(buffer[4:], obj, pos + 4)
        elif obj['pppoe']['type'] == 0x57:
            self.ip6(buffer[4:], obj, pos + 4)
        else:
            print('Unknown ppp.type', obj)

    def mpls(self, buffer, obj, pos):
        offset = 0
        while offset + 5 < len(buffer):
            S = buffer[offset + 2] & 0x1
            offset += 4
            if S:
                mpls_type = buffer[offset] >> 4
                if mpls_type == 4:
                    self.ip4(buffer[offset:], obj, pos + offset)
                    return
                elif mpls_type == 6:
                    self.ip6(buffer[offset:], obj, pos + offset)
                    return
                else:
                    print('Unknown mpls.type', mpls_type, obj, offset)
                    return

    def ethertyperun(self, type, buffer, obj, pos):
        if type == 0x0800:
            self.ip4(buffer, obj, pos)
        elif type == 0x0806:  # ARP
            obj.setdefault('ether', {})['data'] = buffer
        elif type == 0x86dd:
            self.ip6(buffer, obj, pos)
        elif type == 0x8864:
            self.pppoe(buffer, obj, pos)
        elif type == 0x8847:
            self.mpls(buffer, obj, pos)
        elif type == 0x6558:
            self.ether(buffer, obj, pos)
        elif type == 0x6559:
            self.framerelay(buffer, obj, pos)
        elif type == 0x880b:
            self.ppp(buffer, obj, pos)
        else:
            return False
        return True

    def ethertype(self, buffer, obj, pos):
        obj.setdefault('ether', {})['type'] = int.from_bytes(buffer[0:2], byteorder='big')

        if self.ethertyperun(obj['ether']['type'], buffer[2:], obj, pos + 2):
            return

        if obj['ether']['type'] in {0x8100, 0x88a8}:  # VLAN or Q-in-Q
            self.ethertype(buffer[4:], obj, pos + 4)
        else:
            obj['ether']['data'] = buffer[2:]
            # print("Unknown ether.type", obj)

    def ether(self, buffer, obj, pos):
        obj['ether'] = {
            'length': len(buffer),
            'addr1': buffer[0:6].hex(),
            'addr2': buffer[6:12].hex()
        }
        self.ethertype(buffer[12:], obj, pos + 12)

    def radiotap(self, buffer, obj, pos):
        l = buffer[2] + 24 + 6
        ethertype = int.from_bytes(buffer[l:l + 2], byteorder='big')

        if self.ethertyperun(ethertype, buffer[l + 2:], obj, pos + l + 2):
            return

    def nflog(self, buffer, obj, pos):
        offset = 4
        while offset + 4 < len(buffer):
            length = int.from_bytes(buffer[offset:offset + 2], byteorder='little')
            if buffer[offset + 3] == 0 and buffer[offset + 2] == 9:
                if buffer[0] == 2:
                    return self.ip4(buffer[offset + 4:], obj, pos + offset + 4)
                else:
                    return self.ip6(buffer[offset + 4:], obj, pos + offset + 4)
            else:
                offset += ((length + 3) & 0xfffc)

        l = buffer[2] + 24
        if buffer[l + 6] == 0x08 and buffer[l + 7] == 0x00:
            self.ip4(buffer[l + 8:], obj, pos + l + 8)
        elif buffer[l + 6] == 0x86 and buffer[l + 7] == 0xdd:
            self.ip6(buffer[l + 8:], obj, pos + l + 8)

    def framerelay(self, buffer, obj, pos):
        if buffer[2] == 0x03 or buffer[3] == 0xcc:
            self.ip4(buffer[4:], obj, pos + 4)
        elif buffer[2] == 0x08 or buffer[3] == 0x00:
            self.ip4(buffer[4:], obj, pos + 4)
        elif buffer[2] == 0x86 or buffer[3] == 0xdd:
            self.ip6(buffer[4:], obj, pos + 4)

    def pcap__(self, buffer, obj):
        if self.bigEndian:
            obj['pcap'] = {
                'ts_sec': int.from_bytes(buffer[0:4], byteorder='big'),
                'ts_usec': int.from_bytes(buffer[4:8], byteorder='big'),
                'incl_len': int.from_bytes(buffer[8:12], byteorder='big'),
                'orig_len': int.from_bytes(buffer[12:16], byteorder='big')
            }
        else:
            obj['pcap'] = {
                'ts_sec': int.from_bytes(buffer[0:4], byteorder='little'),
                'ts_usec': int.from_bytes(buffer[4:8], byteorder='little'),
                'incl_len': int.from_bytes(buffer[8:12], byteorder='little'),
                'orig_len': int.from_bytes(buffer[12:16], byteorder='little')
            }

        if self.nanosecond:
            obj['pcap']['ts_usec'] //= 1000

        link_type = self.linkType
        incl_len = obj['pcap']['incl_len']
        offset = 16

        if link_type == 0:  # NULL
            if buffer[16] == 30:
                self.ip6(buffer[20:incl_len + 16], obj, 20)
            else:
                self.ip4(buffer[20:incl_len + 16], obj, 20)
        elif link_type == 1:  # Ether
            self.ether(buffer[16:incl_len + 16], obj, 16)
        elif link_type in {12, 101}:  # LOOP or RAW
            if (buffer[16] & 0xF0) == 0x60:
                self.ip6(buffer[16:incl_len + 16], obj, 16)
            else:
                self.ip4(buffer[16:incl_len + 16], obj, 16)
        elif link_type == 107:  # Frame Relay
            self.framerelay(buffer[16:incl_len + 16], obj, 16)
        elif link_type == 113:  # SLL
            obj['ether'] = {}
            self.ethertype(buffer[30:incl_len + 16], obj, 30)
        elif link_type == 127:  # radiotap
            self.radiotap(buffer[16:incl_len + 16], obj, 16)
        elif link_type == 228:  # RAW
            self.ip4(buffer[16:incl_len + 16], obj, 16)
        elif link_type == 239:  # NFLOG
            self.nflog(buffer[16:incl_len + 16], obj, 16)
        elif link_type == 276:  # SLL2
            self.ip4(buffer[36:incl_len + 20], obj, 36)
        else:
            print('Unsupported pcap file', self.filename, 'link type', link_type)

    def decode(self, buffer, obj):
        self.read_header()
        self.pcap__(buffer, obj)
